// Copyright (c) 2022 ChenJun
// Licensed under the MIT License.

// OpenCV
#include <opencv2/core.hpp>
#include <opencv2/core/base.hpp>
#include <opencv2/core/mat.hpp>
#include <opencv2/core/types.hpp>
#include <opencv2/imgproc.hpp>

// STD
#include <algorithm>
#include <cmath>
#include <vector>

#include "armor_detector/detector.hpp"
#include "auto_aim_interfaces/msg/debug_armor.hpp"
#include "auto_aim_interfaces/msg/debug_light.hpp"

namespace rm_auto_aim
{
Detector::Detector(
  const int & bin_thres, const int & color, const LightParams & l, const ArmorParams & a)
: binary_thres(bin_thres), detect_color(color), l(l), a(a)
{
}

std::vector<Armor> Detector::detect(const cv::Mat & input)
{
  binary_img = preprocessImage(input);
  lights_ = findLights(input, binary_img);
  armors_ = matchLights(lights_);

  if (!armors_.empty()) {
    classifier->extractNumbers(input, armors_);
    classifier->classify(armors_);
  }

  return armors_;
}
/*将彩色图转化为灰度图*/
cv::Mat Detector::preprocessImage(const cv::Mat & rgb_img)
{
  cv::Mat gray_img;
  // 将RGB图像转换为灰度图像
  cv::cvtColor(rgb_img, gray_img, cv::COLOR_RGB2GRAY);//gray_img是输出
  cv::Mat binary_img;
  // 对灰度图像进行二值化处理
  cv::threshold(gray_img, binary_img, binary_thres, 255, cv::THRESH_BINARY);//binary_img是输出
  // 返回二值化图像
  return binary_img;
}

std::vector<Light> Detector::findLights(const cv::Mat& rbg_img, const cv::Mat& binary_img)
{
    using std::vector;
    vector<vector<cv::Point>> contours;
    vector<cv::Vec4i> hierarchy;//存储轮廓的层级信息，表示其与父轮廓、子轮廓、下一个轮廓和前一个轮廓的关系
    // 寻找二值化图像中的轮廓
    cv::findContours(binary_img, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
    //cv::RETR_EXTERNAL轮廓检索模式，它告诉 cv::findContours 仅检索最外层的轮廓；cv::CHAIN_APPROX_SIMPLE: 这是轮廓近似方法，它告诉 cv::findContours 如何简化轮廓的点
    vector<Light> lights;
    this->debug_lights.data.clear(); // 清空调试信息
    // 遍历每个轮廓
    for (const auto& contour : contours) {
        // 如果轮廓的点数小于5，则跳过
        if (contour.size() < 5) continue;
        // 计算轮廓的最小外接矩形
        auto r_rect = cv::minAreaRect(contour);
        // 使用灯条矩阵作为参数构造 Light 对象
        auto light = Light(r_rect);
        // 判断是否为有效的灯条
        if (isLight(light)) {
            auto rect = light.boundingRect();
            // 避免断言失败；rect.x + rect.width <= rbg_img.cols、rect.y + rect.height <= rbg_img.rows是确保rect 是否完全位于原始图像 rbg_img 的边界内。
            if (0 <= rect.x && 0 <= rect.width && rect.x + rect.width <= rbg_img.cols && 0 <= rect.y &&
                0 <= rect.height && rect.y + rect.height <= rbg_img.rows) {
                int sum_r = 0, sum_b = 0;
                auto roi = rbg_img(rect); //按照外接矩形从彩色图片中抠出对应的部分
                // 遍历ROI中的每个像素；rows是行数，cols是列数
                for (int i = 0; i < roi.rows; i++) {
                    for (int j = 0; j < roi.cols; j++) {
                      //判断像素点 (j + rect.x, i + rect.y) 是否位于之前计算出的轮廓 contour 内部
                        if (cv::pointPolygonTest(contour, cv::Point2f(j + rect.x, i + rect.y), false) >= 0) {
                            // 如果点在轮廓内
                            sum_r += roi.at<cv::Vec3b>(i, j)[0];//累加红色通道的像素值
                            sum_b += roi.at<cv::Vec3b>(i, j)[2];//累加蓝色通道的像素值
                        }
                    }
                }
                // 根据红蓝像素之和判断颜色--若蓝色通道值大于红色，则判断灯条为蓝色。
                light.color = sum_r > sum_b ? RED : BLUE;
                // 将灯条加入容器中
                lights.emplace_back(light);
            }
        }
    }
    // 返回灯条容器
    return lights;
}
/*通过计算疑似灯条的物体的比例和角度来判断该物体是否是灯条*/
bool Detector::isLight(const Light & light)
{
  // The ratio of light (short side / long side)
  float ratio = light.width / light.length;
  bool ratio_ok = l.min_ratio < ratio && ratio < l.max_ratio;

  bool angle_ok = light.tilt_angle < l.max_angle;

  bool is_light = ratio_ok && angle_ok;

  // Fill in debug information
  auto_aim_interfaces::msg::DebugLight light_data;
  light_data.center_x = light.center.x;
  light_data.ratio = ratio;
  light_data.angle = light.tilt_angle;
  light_data.is_light = is_light;
  this->debug_lights.data.emplace_back(light_data);

  return is_light;
}

//将装有 有效灯条 的容器作为输入参数
std::vector<Armor> Detector::matchLights(const std::vector<Light>& lights)
{
    std::vector<Armor> armors;
    this->debug_armors.data.clear(); // 清空调试信息
    // 遍历所有灯条的组合
    for (auto light_1 = lights.begin(); light_1 != lights.end(); light_1++) {
        for (auto light_2 = light_1 + 1; light_2 != lights.end(); light_2++) {
            // 如果灯条的颜色不符合检测要求，跳过
            if (light_1->color != detect_color || light_2->color != detect_color) continue;
            // 如果灯条包含在其他装甲板内（我们找灯条就是为了确定装甲板，已经组成装甲板的灯条不再使用），跳过
            if (containLight(*light_1, *light_2, lights)) {
                continue;
            }
            // 判断两个灯条是否形成装甲板
            auto type = isArmor(*light_1, *light_2);
            if (type != ArmorType::INVALID) {
                // 用两个符合条件的灯条来创建装甲板对象
                auto armor = Armor(*light_1, *light_2);
                armor.type = type;
                // 将装甲板加入容器中
                armors.emplace_back(armor);
            }
        }
    }
    // 返回装甲板容器
    return armors;
}

// Check if there is another light in the boundingRect formed by the 2 lights
bool Detector::containLight(
  const Light & light_1, const Light & light_2, const std::vector<Light> & lights)
{
  auto points = std::vector<cv::Point2f>{light_1.top, light_1.bottom, light_2.top, light_2.bottom};
  auto bounding_rect = cv::boundingRect(points);

  for (const auto & test_light : lights) {
    if (test_light.center == light_1.center || test_light.center == light_2.center) continue;

    if (
      bounding_rect.contains(test_light.top) || bounding_rect.contains(test_light.bottom) ||
      bounding_rect.contains(test_light.center)) {
      return true;
    }
  }

  return false;
}
/*根据灯条的物理特征来判断它们是否可以组成一个装甲板*/
ArmorType Detector::isArmor(const Light & light_1, const Light & light_2)
{
  //1.长度比检查Ratio of the length of 2 lights (short side / long side)
  float light_length_ratio = light_1.length < light_2.length ? light_1.length / light_2.length
                                                             : light_2.length / light_1.length;
  bool light_ratio_ok = light_length_ratio > a.min_light_ratio;

  //2.长度比检查Distance between the center of 2 lights (unit : light length)
  /*计算两个灯条中心点的欧氏距离，然后将其与预设的最小和最大中心点距离阈值（a.min_small_center_distance 和 a.max_small_center_distance
   或 a.min_large_center_distance 和 a.max_large_center_distance）进行比较，以确定这个距离是否在可接受的范围内。*/
  float avg_light_length = (light_1.length + light_2.length) / 2;
  float center_distance = cv::norm(light_1.center - light_2.center) / avg_light_length;
  bool center_distance_ok = (a.min_small_center_distance <= center_distance &&
                             center_distance < a.max_small_center_distance) ||
                            (a.min_large_center_distance <= center_distance &&
                             center_distance < a.max_large_center_distance);

  //3.角度检查Angle of light center connection计算连接两个灯条中心点的向量，并计算这个向量与x轴正方向的夹角。然后检查这个角度是否小于预设的最大角度阈值（a.max_angle）
  cv::Point2f diff = light_1.center - light_2.center;
  float angle = std::abs(std::atan(diff.y / diff.x)) / CV_PI * 180;
  bool angle_ok = angle < a.max_angle;

  bool is_armor = light_ratio_ok && center_distance_ok && angle_ok;

  // Judge armor type判断装甲板类型（LARGE 或 SMALL）。
  ArmorType type;
  if (is_armor) {
    type = center_distance > a.min_large_center_distance ? ArmorType::LARGE : ArmorType::SMALL;
  } else {
    type = ArmorType::INVALID;
  }

  // Fill in debug information填充了调试信息，包括装甲板类型、灯条长度比、中心点距离和角度，这些信息可以用于后续的调试和分析。
  auto_aim_interfaces::msg::DebugArmor armor_data;
  armor_data.type = ARMOR_TYPE_STR[static_cast<int>(type)];
  armor_data.center_x = (light_1.center.x + light_2.center.x) / 2;
  armor_data.light_ratio = light_length_ratio;
  armor_data.center_distance = center_distance;
  armor_data.angle = angle;
  this->debug_armors.data.emplace_back(armor_data);

  return type;
}
/*在detector_node函数中使用*/
/*将装甲板中的所有数字图像整合到一个单独的图像中，以便进一步对数字进行分析和识别。通过合并所有数字图像，可以方便地从装甲板图像中提取出数字，并进行数字识别或其他相关操作*/
cv::Mat Detector::getAllNumbersImage()
{
  //首先，检查 armors_（装甲板对象容器）是否为空。
  if (armors_.empty()) {
    //如果 armors_ 为空，表示没有装甲板对象，那么函数会创建一个空白的图像，大小为 20x28，像素值为 0。
    return cv::Mat(cv::Size(20, 28), CV_8UC1);
  } 
  else {
    //如果 armors_ 不为空，表示存在装甲板对象，那么函数会遍历每个装甲板对象。
    std::vector<cv::Mat> number_imgs;
    number_imgs.reserve(armors_.size());
    for (auto & armor : armors_) {
      //对于每个装甲板对象，会获取其中的数字图像（number_img），并将其存储在 number_imgs 容器中。
      number_imgs.emplace_back(armor.number_img);
    }
    cv::Mat all_num_img;
    //最后，函数使用 cv::vconcat 函数将 number_imgs 容器中的图像垂直合并为一个大的图像，并将合并后的图像返回
    cv::vconcat(number_imgs, all_num_img);
    return all_num_img;
  }
}

void Detector::drawResults(cv::Mat & img)
{
  // Draw Lights
  for (const auto & light : lights_) {
    cv::circle(img, light.top, 3, cv::Scalar(255, 255, 255), 1);
    cv::circle(img, light.bottom, 3, cv::Scalar(255, 255, 255), 1);
    auto line_color = light.color == RED ? cv::Scalar(255, 255, 0) : cv::Scalar(255, 0, 255);
    cv::line(img, light.top, light.bottom, line_color, 1);
  }

  // Draw armors
  for (const auto & armor : armors_) {
    cv::line(img, armor.left_light.top, armor.right_light.bottom, cv::Scalar(0, 255, 0), 2);
    cv::line(img, armor.left_light.bottom, armor.right_light.top, cv::Scalar(0, 255, 0), 2);
  }

  // Show numbers and confidence显示数字和置信度信息
  for (const auto & armor : armors_) {
    cv::putText(
      img, armor.classfication_result, armor.left_light.top, cv::FONT_HERSHEY_SIMPLEX, 0.8,
      cv::Scalar(0, 255, 255), 2);
  }
}

}  // namespace rm_auto_aim
